import datetime
import numpy as np
import rfi
import foregrounds

class Simulator:
    """
    Class: Simulator generates radio data, including rfi
    Args:
        freqs: frequency 
        times: time 
        data_filename: filename of data to be saved
    """
    def __init__(self,freqs,times,data_filename=None):
        if freqs is None:
            raise ValueError("freqs must be given")
        if times is None:
            raise ValueError("times must be given")
        self.freqs = freqs
        self.times = times
        self.rfi_data = np.zeros((times.size, freqs.size), dtype=np.complex128)
        self.rfi_flag = np.zeros((times.size, freqs.size), dtype=np.int16)
        self.RFI = rfi.Rfi(self.rfi_data,self.rfi_flag,self.freqs,self.times)
        self.data_array = np.zeros((times.size, freqs.size), dtype=np.complex128)
        self.data_mean = 0 # Average of background amplitude magnitude
        if data_filename is None:
            # if is none, replace it with current time
            self.data_filename = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
 
    def add_foregrounds(self, model=None):
        """
        Add a foreground model to the visibilities.
        Args:
            model: Calling different types of functions
        """
        if model is None:
            raise ValueError("model must be given")
        fgrd,data_mean = getattr(foregrounds,model)(self.times,self.freqs);
        self.data_array += fgrd
        self.data_mean = data_mean

    def add_rfi(self,model=None,**kwargs):
        """
        Add a rfi model to the visibilities.
        Args:
            model: Calling different types of functions
        """
        if model is None:
            raise ValueError("model must be given")
        rfid,rfif = getattr(self.RFI,model)(data_mean=self.data_mean,**kwargs)
        self.data_array += rfid
        self.rfi_flag = rfif


    def add_noise(self,model):
        """
        Add thermal noise model to the visibilities.
        Args:
            model: Calling different types of functions
        """
